import os
import json
import logging
import asyncio
import aiohttp

from dotenv import load_dotenv
from pathlib import Path
from typing import List, Annotated, Optional, Literal, Sequence, TypedDict, Union
from pydantic import BaseModel, Field
from langchain_deepseek import ChatDeepSeek
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.prompts import PromptTemplate
from langchain_mcp_adapters.client import MultiServerMCPClient
from langgraph.checkpoint.memory import InMemorySaver
from langchain_core.tools import tool
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage, ToolMessage, BaseMessage
from langchain_core.messages.ai import AIMessageChunk, ToolCallChunk
from langchain_core.messages.tool import ToolMessage
from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode, tools_condition
from langchain_core.runnables import RunnableConfig

from insight_agent.mcp_client.frontend_service.configuration import Configuration


load_dotenv(override=True)
cfg = Configuration()
thread_config_1 = {"configurable": {"thread_id": "thread1"}}
checkpointer = InMemorySaver()

# ------------------------

# 1. 定义状态（State）
# ------------------------
class AgentState(TypedDict):
    messages: Annotated[Sequence, add_messages]
    # 标记是否正在流式输出
    streaming: bool


# ------------------------
# 2. 创建异步工具（Async Tools）
# ------------------------
@tool
async def get_weather(city: str) -> str:
    """根据城市名获取当前天气信息。模拟异步API调用。"""
    # 模拟异步网络请求
    print('\n开始调用 get_weather ...', end='')
    await asyncio.sleep(0.5)  # 模拟网络延迟
    weather_data = {
        "北京": "晴，25°C",
        "上海": "多云，23°C",
        "深圳": "阵雨，28°C"
    }
    print('\n结束调用 get_weather ...', end='')
    return weather_data.get(city.lower(), f"抱歉，找不到{city}的天气信息。")


@tool
async def get_stock_price(company_name: str) -> str:
    """根据公司名获取当前股价。模拟异步API调用。"""
    # 模拟异步网络请求
    print('\n开始调用 get_stock_price ...', end='')
    await asyncio.sleep(1.5)  # 模拟更长的网络延迟
    stock_data = {
        "苹果公司": "$172.35",
        "微软公司": "$412.22",
        "谷歌公司": "$156.78"
    }
    print('\n结束调用 get_stock_price ...', end='')
    return stock_data.get(company_name.lower(), f"抱歉，找不到{company_name}的股价信息。")


@tool
async def search_web(query: str) -> str:
    """搜索网络获取最新信息。模拟异步API调用。"""
    print('\n开始调用 search_web ...', end='')
    await asyncio.sleep(1)  # 模拟网络延迟
    # 这里是模拟数据，实际应用中会调用真正的搜索API
    print('\n结束调用 search_web ...', end='')
    return f"关于'{query}'的搜索结果: 最新的相关信息摘要为123456"



async def load_tools(servers_config_path):
    with open(servers_config_path, "r", encoding="utf-8") as f:
        mcp_servers = json.load(f).get("mcp_servers", {})
    mcp_client = MultiServerMCPClient(mcp_servers)
    tools = await mcp_client.get_tools() + [get_weather, get_stock_price, search_web]
    tool_map = {tool.name: tool for tool in tools}
    return tools, tool_map

# 将所有工具放在一个列表中
tools, tool_map = asyncio.run(load_tools(f'../{os.getenv("MCP_SERVERS_CONFIG_PATH")}'))


# ------------------------
# 3. 设置模型和绑定工具
# ------------------------
# 初始化LLM - 使用支持并行工具调用的模型
model = ChatDeepSeek(model='deepseek-chat')
prompt_path = f'../{os.getenv("MAIN_AGENT_PROMPT_PATH")}'
with open(prompt_path, "r", encoding="utf-8") as f:
    prompt = f.read()

# 将提示词和工具绑定到模型
model_with_tools = model.bind_tools(tools)


# ------------------------
# 4. 定义图的节点（Nodes）
# ------------------------
async def call_model(state: AgentState, config: Optional[RunnableConfig] = None):
    """Agent节点：调用LLM，决定是否调用工具及调用哪些工具。"""
    response = await model_with_tools.ainvoke(state["messages"], config=config)
    return {"messages": [response]}


async def call_tools(state: AgentState):
    """
    Tools节点：并行、异步执行所有被调用的工具。
    支持流式输出工具执行进度。
    """
    last_message = state["messages"][-1]
    print('state: ', state)
    print('last_message: ', last_message)

    if not hasattr(last_message, "tool_calls") or not last_message.tool_calls:
        return {"messages": []}

    # 并行执行所有工具调用
    tool_coroutines = []
    for tool_call in last_message.tool_calls:
        if tool_call["name"] in tool_map:
            tool_coroutines.append(
                tool_map[tool_call["name"]].ainvoke(tool_call["args"])
            )
        else:
            # 如果工具不存在，返回错误信息
            tool_coroutines[-1].result = f"错误：工具 '{tool_call['name']}' 未找到。"

    # 使用asyncio.gather并行执行所有工具
    tool_results = await asyncio.gather(*tool_coroutines, return_exceptions=True)

    # 处理结果
    tool_messages = []
    for i, (tool_call, result) in enumerate(zip(last_message.tool_calls, tool_results)):
        if isinstance(result, Exception):
            # 处理异常
            tool_messages.append(ToolMessage(
                content=f"调用工具 {tool_call['name']} 时出错: {str(result)}",
                name=tool_call['name'],
                tool_call_id=tool_call['id']
            ))
        else:
            # 正常结果
            tool_messages.append(ToolMessage(
                content=str(result),
                name=tool_call['name'],
                tool_call_id=tool_call['id']
            ))

    return {"messages": tool_messages}


# ------------------------
# 5. 流式输出处理
# ('messages', (AIMessageChunk(content='', additional_kwargs={'tool_calls': [{'index': 0, 'id': 'call_0_6311d369-b1d4-4cde-b5b3-3fcfeb003373', 'function': {'arguments': '', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={}, id='run--17a6509e-5c43-439c-9577-0a7150a3e323', tool_calls=[{'name': 'get_weather', 'args': {}, 'id': 'call_0_6311d369-b1d4-4cde-b5b3-3fcfeb003373', 'type': 'tool_call'}], tool_call_chunks=[{'name': 'get_weather', 'args': '', 'id': 'call_0_6311d369-b1d4-4cde-b5b3-3fcfeb003373', 'index': 0, 'type': 'tool_call_chunk'}]), {'langgraph_step': 1, 'langgraph_node': 'agent', 'langgraph_triggers': ('branch:to:agent',), 'langgraph_path': ('__pregel_pull', 'agent'), 'langgraph_checkpoint_ns': 'agent:012bb2dc-3864-55e7-b55b-393bdebff9dc', 'checkpoint_ns': 'agent:012bb2dc-3864-55e7-b55b-393bdebff9dc', 'ls_provider': 'openai', 'ls_model_name': 'deepseek-chat', 'ls_model_type': 'chat', 'ls_temperature': None}))
# ('updates', {'agent': {'messages': [AIMessage(content='我来帮您查询上海的天气情况。', additional_kwargs={'tool_calls': [{'index': 0, 'id': 'call_0_6311d369-b1d4-4cde-b5b3-3fcfeb003373', 'function': {'arguments': '{"city": "\\u4e0a\\u6d77"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'finish_reason': 'tool_calls', 'model_name': 'deepseek-chat', 'system_fingerprint': 'fp_feb633d1f5_prod0820_fp8_kvcache'}, id='run--17a6509e-5c43-439c-9577-0a7150a3e323', tool_calls=[{'name': 'get_weather', 'args': {'city': '上海'}, 'id': 'call_0_6311d369-b1d4-4cde-b5b3-3fcfeb003373', 'type': 'tool_call'}], usage_metadata={'input_tokens': 255, 'output_tokens': 30, 'total_tokens': 285, 'input_token_details': {'cache_read': 192}, 'output_token_details': {}})]}})
# ------------------------


# ------------------------
# 6. 构建图（Graph）
# ------------------------
# 创建图
graph_builder = StateGraph(AgentState)

# 添加节点
graph_builder.add_node("agent", call_model)
graph_builder.add_node("tools", call_tools)

# 设置入口点
graph_builder.set_entry_point("agent")

# 使用LangGraph内置的条件判断
graph_builder.add_conditional_edges(
    "agent",
    # 使用内置工具条件判断
    tools_condition
)

# 从tools节点出来后，回到agent节点
graph_builder.add_edge("tools", "agent")


# 编译图
main_agent = graph_builder.compile(checkpointer)


# ------------------------
# 7. 测试函数
# ------------------------
async def run_chat_loop() -> None:
    """启动 MCP-Agent 聊天循环"""

    # CLI 多轮对话
    print("\n🤖 MCP Agent 已启动，输入 'quit' 退出")
    while True:
        user_input = input("\n你: ").strip()
        if user_input.lower() == "quit":
            break
        try:
            print("\n🤖 AI:")
            async for chunk in main_agent.astream({"messages":[{"role": "system", "content": prompt}, {"role": "user", "content": user_input}]},
                                                  cfg.thread_config_1,
                                                  stream_mode=['messages', 'updates']):
                # print('👣', chunk)
                if chunk[0] == 'messages':
                    if isinstance(chunk[1][0], AIMessageChunk):
                        response_content = chunk[1][0].content
                        print(response_content, end='')
                        if chunk[1][0].tool_calls:
                            if not chunk[1][0].tool_calls[0]['name']:
                                continue
                            print(f"\n🤖 根据用户输入的信息，Agent决定调用工具: {[tc['name'] for tc in chunk[1][0].tool_calls]}", end='')
                            # print('tool_calls: ', chunk[1][0].tool_calls)
                            # print('tool_call_chunks: ', chunk[1][0].tool_call_chunks)

                elif chunk[0] == 'updates':
                    pass

        except Exception as exc:
            print(f"\n⚠️  出错: {exc}")

        # if step.get("streaming", False):
        #     for chunk in step.get("stream_buffer", []):
        #         print(chunk, end="", flush=True)
        #     if "stream_buffer" in step:
        #         step["stream_buffer"] = []  # 清空缓冲区
        #
        # 查询北京天气和异网用户在地图上的分布
        # 北京市异网用户的分布及疑似诈骗的案例
        # 我想使用图数据库异网拉新，同时查看北京市异网用户的分布
        # 识别北京潜在诈骗的案例，同时查看交往圈和网络圈的个数在地图上的分布
        # 根据个人客户的产品订购情况发现潜在的合作企业，同时查看流量使用超过1GB的客户在地图上的分布



if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
    asyncio.run(run_chat_loop())

"""
## 关键特性说明

### 1. 智能工具调用决策
 Agent 使用大模型自动判断是否需要调用工具：
- 简单问候 → 不调用工具
- 单一问题 → 调用一个相关工具
- 复杂问题 → 并行调用多个工具

### 2. 并行异步执行
- 使用 `asyncio.gather()` 实现真正的并行工具调用
- 所有工具都是异步函数，使用 `async/await` 语法

### 3. 流式输出
- 逐步显示 Agent 的决策过程和工具调用进度
- 实时显示部分结果，而不是等待所有完成

### 4. 错误处理
- 包含完整的异常处理机制
- 当工具不存在或调用失败时提供的错误信息

## TODO:
1. **添加工具调用限制**：添加逻辑限制单次调用工具的数量，防止过度调用
2. **实现工具优先级**：为工具添加优先级，确保关键工具先执行
3. **添加缓存机制**：对频繁使用的工具结果进行缓存，提高响应速度
4. **实现更复杂的流式输出**：使用 HTTP Streaming Response 实现真正的实时流式输出
5. **添加验证机制**：对工具返回结果进行验证，确保数据质量
"""


# # 思维链
#
# # 构建图
# graph = StateGraph(agent)
# graph.set_entry_point("agent")
#
# # 运行并获取思维过程
# async def run_with_thinking(input_text: str):
#     thinking_handler.thinking_steps.clear()  # 清空之前的记录
#
#     result = await graph.ainvoke(
#         {"messages": [("user", input_text)]},
#         config={"callbacks": [thinking_handler]}  # 确保回调生效
#     )
#
#     # 获取完整的思维过程
#     thinking_process = thinking_handler.get_thinking_process()
#
#     return result, thinking_process
#
#
# from langchain.agents import AgentExecutor
# from langchain.agents.format_scratchpad import format_log_to_str
# from langchain.agents.output_parsers import ReActSingleInputOutputParser
#
# # 配置 agent 以输出详细推理过程
# agent = (
#     {
#         "input": lambda x: x["input"],
#         "agent_scratchpad": lambda x: format_log_to_str(x["intermediate_steps"])
#     }
#     | prompt
#     | llm
#     | ReActSingleInputOutputParser()
# )
#
# # 执行并捕获中间步骤
# agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
#
# result = agent_executor.invoke(
#     {"input": "今天的天气怎么样？"},
#     return_intermediate_steps=True  # 关键参数
# )
#
# # 获取思维过程
# intermediate_steps = result["intermediate_steps"]
# for step in intermediate_steps:
#     print(f"动作: {step[0]}")
#     print(f"观察: {step[1]}")
#     print("---")

# 通过 RunnableLambda 动态生成 prompt
# from langchain_core.runnables import RunnableLambda
#
# def prompt_router(state: AgentState):
#     if "technical" in state["input"]:
#         return "技术专家提示：..."
#     else:
#         return "通用助手提示：..."
#
# # 在StateGraph构建时组合
# graph = StateGraph(AgentState)
# graph.add_node("agent",
#     prompt_router | model.bind_tools(tools)  # 动态prompt+模型
# )
